{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3b668138-98ed-41dc-b17e-8f2d643d3e1a",
   "metadata": {},
   "source": [
    "# 0. Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "293222c2-ea98-44d7-808d-06144b1ac015",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install wfdb\n",
    "\n",
    "# Handle data path\n",
    "from pathlib import Path\n",
    "\n",
    "# Read and display data from Physionet\n",
    "import wfdb\n",
    "import pprint\n",
    "import collections\n",
    "from IPython.display import clear_output\n",
    "\n",
    "# Data manipulation and plotting\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.signal import resample\n",
    "from scipy.signal import butter,filtfilt\n",
    "\n",
    "# Divide data into train and test set and save to HDF5\n",
    "import h5py\n",
    "import os\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn import preprocessing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "972f69db-5a5d-42dc-bd18-6e0878d2234d",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_data_path = Path('./data/mit-bih-arrhythmia-database-1.0.0')\n",
    "# main directory that store data\n",
    "data_dir = Path('./data')\n",
    "os.makedirs(data_dir, exist_ok=True)\n",
    "\n",
    "# directory that store original MIT-BIH data\n",
    "\n",
    "mit_arrh_dir = raw_data_path\n",
    "os.makedirs(mit_arrh_dir, exist_ok=True)\n",
    "\n",
    "# directory that store processed data\n",
    "process_dir = data_dir / 'processed_data'\n",
    "temp_dir = process_dir / 'temp'\n",
    "os.makedirs(process_dir, exist_ok=True)\n",
    "\n",
    "raw_data_dir = process_dir\n",
    "os.makedirs(raw_data_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "715d28db-9c11-4d64-93d0-f39abb704b43",
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate_beat_type(beat_type_arr, peak_loc_arr):\n",
    "    n_type = ['N', 'L', 'R', 'e', 'j']\n",
    "    s_type = ['A', 'a', 'J', 'S']\n",
    "    v_type = ['V', 'E']\n",
    "    f_type = ['F']\n",
    "    q_type = ['/', 'f', 'Q']\n",
    "    \n",
    "    not_used = ['B', 'r', 'n', '?']\n",
    "    \n",
    "    encoded_label = []\n",
    "    peak_loc_label = []\n",
    "    ignored_label = []\n",
    "    \n",
    "    for beat_type, peak_loc in zip(beat_type_arr, peak_loc_arr):\n",
    "        if beat_type in n_type:\n",
    "            encoded_label.append(0)\n",
    "        elif beat_type in s_type:\n",
    "            encoded_label.append(1)\n",
    "        elif beat_type in v_type:\n",
    "            encoded_label.append(2)\n",
    "        elif beat_type in f_type:\n",
    "            encoded_label.append(3)\n",
    "        elif beat_type in q_type:\n",
    "            encoded_label.append(4)\n",
    "        else:\n",
    "            ignored_label.append(beat_type)\n",
    "            continue\n",
    "        peak_loc_label.append(peak_loc)\n",
    "        \n",
    "        \n",
    "            \n",
    "    encoded_label = np.array(encoded_label)\n",
    "    peak_loc_label = np.array(peak_loc_label)\n",
    "\n",
    "    return encoded_label, peak_loc_label, ignored_label\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a74e550-b381-4ba9-913a-3be68061eadf",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Noted: This part of data is done incorrectly. The dataset should be split in train and test using records as follow instead of random split. The reason for this is from the paper:\n",
    "## How the Choice of Samples for Building Arrhythmia Classifiers ImpactTheir Performances\n",
    "## Train dataset record number: 101, 106, 108, 109, 112, 114, 115,116, 118, 119, 122, 124, 201, 203, 205, 207, 208, 209, 215,220, 223, and 230,\n",
    "## Test dataset record number: 100, 103, 105, 11, 113, 117, 121, 123, 200, 202, 210, 212, 213,214, 219, 221, 222, 228, 231, 232, and 234\n",
    "## read data and annotation\n",
    "\n",
    "file = open(mit_arrh_dir/'RECORDS')\n",
    "record_list = file.read().splitlines()\n",
    "file.close()\n",
    "# progress bar to check if the program running\n",
    "def update_progress(progress):\n",
    "    bar_length = 50\n",
    "    if isinstance(progress, int):\n",
    "        progress = float(progress)\n",
    "    if not isinstance(progress, float):\n",
    "        progress = 0\n",
    "    if progress < 0:\n",
    "        progress = 0\n",
    "    if progress >= 1:\n",
    "        progress = 1\n",
    "        \n",
    "    block = int(round(bar_length * progress))\n",
    "\n",
    "    clear_output(wait = True)\n",
    "    text = \"Progress: [{0}] {1:.1f}%\".format( \"#\" * block + \"-\" * (bar_length - block), progress * 100)\n",
    "    print(text)\n",
    "    \n",
    "bad_labels = []\n",
    "with h5py.File(raw_data_dir / 'mit_raw_data.h5', 'w') as file:\n",
    "    for i in range(len(record_list)):\n",
    "        summary = wfdb.rdrecord(str(mit_arrh_dir / record_list[i])).__dict__\n",
    "        \n",
    "        ## only get the second signal which is ii - chest ECG\n",
    "        anno = wfdb.rdann(str(mit_arrh_dir / record_list[i]), extension='atr', summarize_labels=True).__dict__\n",
    "        peaks = anno['sample']\n",
    "        labels = anno['symbol']\n",
    "        encoded_labels, peak_labels, ignored_label = translate_beat_type(labels, peaks)\n",
    "        bad_labels += ignored_label\n",
    "        \n",
    "        file.create_dataset('data/{:03d}'.format(i), data=summary['p_signal'][:, 0])\n",
    "        file.create_dataset('peak/{:03d}'.format(i), data=peak_labels)\n",
    "        file.create_dataset('label/{:03d}'.format(i), data=encoded_labels)\n",
    "        \n",
    "        update_progress((i+1) / float(len(record_list)))\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2923659a-f54a-4ef2-a97f-cc0453d8aaa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "\n",
    "def check_data_group(file):\n",
    "    data_group = file['data']\n",
    "    total = len(data_group.keys())\n",
    "    index = 0\n",
    "    total = []\n",
    "    peaks = []\n",
    "    for key in data_group.keys():\n",
    "        temp_data = list(file['label/{}'.format(key)])\n",
    "        total += temp_data\n",
    "        temp_data = list(file['peak/{}'.format(key)])\n",
    "        peaks += temp_data\n",
    "    \n",
    "    print(np.unique(total))\n",
    "    print(collections.Counter(total))\n",
    "    print(len(total) - collections.Counter(total)[5])\n",
    "    print(len(peaks))\n",
    "    print(len(total))\n",
    "\n",
    "with h5py.File(raw_data_dir / 'mit_raw_data.h5', 'r') as file:\n",
    "    check_data_group(file)\n",
    "    \n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e61b8e99-7e6f-4083-a79f-db2c65102846",
   "metadata": {},
   "outputs": [],
   "source": [
    "def butter_highpass_filter(data, cutoff, fs, order):\n",
    "    nyq = 0.5 * fs\n",
    "    normal_cutoff = cutoff / nyq\n",
    "    # Get the filter coefficients \n",
    "    b, a = butter(order, normal_cutoff, btype='high', analog=False)\n",
    "    y = filtfilt(b, a, data)\n",
    "    return y\n",
    "\n",
    "# progress bar to check if the program running\n",
    "def update_progress(progress):\n",
    "    bar_length = 50\n",
    "    if isinstance(progress, int):\n",
    "        progress = float(progress)\n",
    "    if not isinstance(progress, float):\n",
    "        progress = 0\n",
    "    if progress < 0:\n",
    "        progress = 0\n",
    "    if progress >= 1:\n",
    "        progress = 1\n",
    "        \n",
    "    block = int(round(bar_length * progress))\n",
    "\n",
    "    clear_output(wait = True)\n",
    "    text = \"Progress: [{0}] {1:.1f}%\".format( \"#\" * block + \"-\" * (bar_length - block), progress * 100)\n",
    "    print(text)\n",
    "\n",
    "filtered_ecg_data = []\n",
    "ecg_peak_annotations = []\n",
    "\n",
    "with h5py.File(raw_data_dir / 'mit_raw_data.h5', 'r') as file:\n",
    "    ## only get the second signal which is ii - chest ECG\n",
    "    data_group = file['data']\n",
    "    total = len(data_group.keys())\n",
    "    index = 0\n",
    "    for key in data_group.keys():\n",
    "        temp_data = file['data/{}'.format(key)]\n",
    "        filter_data = butter_highpass_filter(temp_data, 0.25, 360, 2)\n",
    "        filtered_ecg_data.append(filter_data)\n",
    "        temp_peak = file['peak/{}'.format(key)][:]\n",
    "        ecg_peak_annotations.append(temp_peak)\n",
    "        update_progress((index+1) / float(total))\n",
    "        index += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bfeee50-b77a-4767-a072-778e0d647d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(arr):\n",
    "    arr_range = np.ptp(arr)\n",
    "    arr_min = np.min(arr)\n",
    "    norm_arr = (arr - arr_min) / float(arr_range)\n",
    "    return norm_arr\n",
    "\n",
    "ecg_data_norm = []\n",
    "for i in range(len(filtered_ecg_data)):\n",
    "    signal = normalize(filtered_ecg_data[i])\n",
    "    ecg_data_norm.append(signal)\n",
    "    update_progress( (i + 1) / float(len(filtered_ecg_data)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8086a121-0ef5-493f-9322-add7562233aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_local_peak_location(beat_locs, ecg_sample, sample_freq):\n",
    "    interval = 10 * sample_freq\n",
    "    max_length = len(ecg_sample)\n",
    "    correct_locs = []\n",
    "    try:\n",
    "        for i in range(len(beat_locs)):\n",
    "            start_base_line = beat_locs[i] - interval / 2\n",
    "            stop_base_line = beat_locs[i] + interval / 2\n",
    "            \n",
    "            # left end of data\n",
    "            if start_base_line < 0 :\n",
    "                stop_base_line += abs(start_base_line)\n",
    "                start_base_line = 0\n",
    "\n",
    "            # right end of data\n",
    "            if stop_base_line >= max_length :\n",
    "                start_base_line -= abs(stop_base_line - max_length)\n",
    "                stop_base_line = max_length\n",
    "\n",
    "            # local mean base line\n",
    "            base_line = np.mean(ecg_sample[int(start_base_line):int(stop_base_line)])\n",
    "            \n",
    "            local_sample = None\n",
    "            \n",
    "            # left end of data\n",
    "            if beat_locs[i] - 2 < 0:\n",
    "                local_sample = ecg_sample[beat_locs[i] : beat_locs[i] + 5]\n",
    "            \n",
    "            # right end of data\n",
    "            elif beat_locs[i] + 3 > max_length:\n",
    "                local_sample = ecg_sample[beat_locs[i] - 5: beat_locs[i]]\n",
    "            \n",
    "            # normal case\n",
    "            else:\n",
    "                local_sample = ecg_sample[beat_locs[i] - 2: beat_locs[i] + 3]\n",
    "                \n",
    "            diff = [abs(x - base_line) for x in local_sample]\n",
    "            max_loc = np.argmax(diff)\n",
    "            new_loc = beat_locs[i] - 2 + max_loc\n",
    "            correct_locs.append(new_loc)\n",
    "    except:\n",
    "        ## error case print stuff\n",
    "        print(beat_locs[i])\n",
    "        print(local_sample)\n",
    "        print(base_line)\n",
    "        print(diff)\n",
    "        print(max_loc)\n",
    "    correct_locs = np.array(correct_locs)\n",
    "    return correct_locs\n",
    "\n",
    "## Apply to the rest of the data\n",
    "ecg_correct_peak_location = []\n",
    "with h5py.File(raw_data_dir / 'mit_raw_data.h5', 'r') as file:\n",
    "    ## only get the second signal which is ii - chest ECG\n",
    "    data_group = file['peak']\n",
    "    total = len(data_group.keys())\n",
    "    index = 0\n",
    "    for key in data_group.keys():\n",
    "        temp_peak_data = file['peak/{}'.format(key)]\n",
    "        corrected_loc = find_local_peak_location(temp_peak_data, ecg_data_norm[int(key)], 360)\n",
    "        ecg_correct_peak_location.append(corrected_loc)\n",
    "        index += 1\n",
    "        update_progress((index)/ float(len(ecg_data_norm)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce0b9b58-c13c-4721-882a-77b30e46da42",
   "metadata": {},
   "outputs": [],
   "source": [
    "## plot another sample\n",
    "\n",
    "sample_num = 14\n",
    "sample_freq = 360\n",
    "ecg_sample = ecg_data_norm[sample_num]\n",
    "time_in_second = np.arange(0, len(ecg_sample)) / sample_freq\n",
    "\n",
    "beat_loc_sample = ecg_correct_peak_location[sample_num]\n",
    "beat_value = np.take(ecg_sample, beat_loc_sample, axis=0)\n",
    "beat_location_in_second = beat_loc_sample / sample_freq\n",
    "\n",
    "plt.plot(time_in_second, ecg_sample)\n",
    "plt.scatter(beat_location_in_second, beat_value, c='r')\n",
    "plt.xlim(1600, 1610)\n",
    "plt.xlabel('time in second')\n",
    "plt.ylabel('ECG value in mV')\n",
    "plt.title('ECG value over time with beat overlay')\n",
    "plt.show()\n",
    "plt.clf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b9f1ccb-33f5-4078-bc6c-a444e9736919",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Extract signal beat by beat\n",
    "\"\"\"\n",
    "    ecg_data - ecg value for all record\n",
    "    ecg_peak_annotations - beat location for all record\n",
    "    beat_type_annotations - beat label for all record\n",
    "\"\"\"\n",
    "\n",
    "## We only looking to classify these 5 type from 0 to 4\n",
    "accepted_labels = [0, 1, 2, 3, 4]\n",
    "\n",
    "def check_for_array_overflow(end_position, max_len):\n",
    "    if end_position > max_len:\n",
    "        end_poistion = max_len\n",
    "    return end_position\n",
    "\n",
    "bad_label = []\n",
    "## extract beat signal starting from peak of the beat to 1.2 * period where period is the interval between beat\n",
    "def extract_beats(ecg_values, beat_locs, beat_labels, accepted_labels):\n",
    "    signals = []\n",
    "    labels = []\n",
    "    max_len = len(ecg_values)\n",
    "    # skip last beat since there is no follow-up beat after it\n",
    "    for i in range(len(beat_locs) - 1):\n",
    "        ## only handle beat with N, V, Q, S, F label\n",
    "        if beat_labels[i] in accepted_labels:\n",
    "            start = beat_locs[i]\n",
    "            ## get 1.2 * interval between beat\n",
    "            period = int((beat_locs[i+1] - beat_locs[i]) * 1.2)\n",
    "            end = start + period\n",
    "            end = check_for_array_overflow(end, max_len)\n",
    "            \n",
    "            signals.append(ecg_values[start: end])\n",
    "            labels.append(beat_labels[i])\n",
    "        else:\n",
    "            bad_label.append(beat_labels[i])\n",
    "    \n",
    "    return signals, labels\n",
    "\n",
    "# progress bar to check if the program running\n",
    "def update_progress(progress):\n",
    "    bar_length = 50\n",
    "    if isinstance(progress, int):\n",
    "        progress = float(progress)\n",
    "    if not isinstance(progress, float):\n",
    "        progress = 0\n",
    "    if progress < 0:\n",
    "        progress = 0\n",
    "    if progress >= 1:\n",
    "        progress = 1\n",
    "        \n",
    "    block = int(round(bar_length * progress))\n",
    "\n",
    "    clear_output(wait = True)\n",
    "    text = \"Progress: [{0}] {1:.1f}%\".format( \"#\" * block + \"-\" * (bar_length - block), progress * 100)\n",
    "    print(text)\n",
    "    \n",
    "ecg_signals = []\n",
    "ecg_labels = []\n",
    "\n",
    "\n",
    "beat_type_annotations = []\n",
    "with h5py.File(raw_data_dir / 'mit_raw_data.h5', 'r') as file:\n",
    "    ## only get the second signal which is ii - chest ECG\n",
    "    data_group = file['label']\n",
    "    total = len(data_group.keys())\n",
    "    index = 0\n",
    "    for key in data_group.keys():\n",
    "        temp_peak_data = file['label/{}'.format(key)][:]\n",
    "        beat_type_annotations.append(temp_peak_data)\n",
    "        index += 1\n",
    "        update_progress((index)/ float(total))\n",
    "        \n",
    "    \n",
    "\n",
    "for i in range(len(ecg_data_norm)):\n",
    "    signal, label = extract_beats(ecg_data_norm[i], ecg_correct_peak_location[i], beat_type_annotations[i], accepted_labels)\n",
    "    ecg_signals.append(signal)\n",
    "    ecg_labels.append(label)\n",
    "    update_progress((i + 1) / float(len(ecg_data_norm)))\n",
    "    \n",
    "## check if the data and label have equal length\n",
    "print('Data length:', len(ecg_signals), 'Label length:', len(ecg_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "befac00f-d9df-4d49-8201-8f74fb202e82",
   "metadata": {},
   "outputs": [],
   "source": [
    "## plot one of the extract beat with length of 1.2 period\n",
    "beat = ecg_signals[0][115]\n",
    "# sampling frequency 360 Hz\n",
    "time_in_seconds = np.arange(len(beat)) / 360.\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"]= 15, 5\n",
    "plt.plot(time_in_seconds, beat)\n",
    "plt.xlabel('time in second')\n",
    "plt.ylabel('ECG value in mV')\n",
    "plt.title('ECG value for one beat')\n",
    "plt.show()\n",
    "plt.clf()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86af1212-18d6-4fcc-b3e0-57b1a217a7fc",
   "metadata": {},
   "source": [
    "# 1. Data split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d89f1eed-d824-41ed-92ea-3f23fb116ca2",
   "metadata": {},
   "source": [
    "## 1.1 Inter-patient split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8e90a8e-35a2-42e4-837e-3fc6d926f9f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# zero pad all heartbeat to 7 seconds\n",
    "import copy\n",
    "ecg_signals_zeropad = copy.deepcopy(ecg_signals)\n",
    "\n",
    "for i in range(len(ecg_signals)):\n",
    "    for j in range(len(ecg_signals[i])):\n",
    "\n",
    "        # pad to 256\n",
    "        zero_pad_max_length = 1024\n",
    "        pad_length = zero_pad_max_length - len(ecg_signals[i][j])\n",
    "        if pad_length <=0:\n",
    "            zero_pad_data = ecg_signals[i][j][:zero_pad_max_length]\n",
    "        else:\n",
    "            zero_pad_data = np.pad(ecg_signals[i][j], (0,pad_length),'constant')\n",
    "        ecg_signals_zeropad[i][j] = zero_pad_data\n",
    "        \n",
    "\n",
    "## plot a zero padded signal\n",
    "ecg_sample = ecg_signals_zeropad[5][100]\n",
    "\n",
    "# sampling frequency 125 Hz\n",
    "time_in_second = np.arange(len(ecg_sample)) / 360.\n",
    "\n",
    "plt.plot(time_in_second, ecg_sample, c='b', label='360 Hz')\n",
    "plt.xlabel('time in second')\n",
    "plt.ylabel('ECG value in mV')\n",
    "plt.title('Zero padded ECG signal for one beat')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "plt.clf()\n",
    "\n",
    "# ecg_signals_zeropad: 48 patients x number_of_beats x 256\n",
    "# ecg_labels: 48 patients x number_of_beats\n",
    "\n",
    "## Split train and test datasets with record number\n",
    "## Train dataset record number: 101, 106, 108, 109, 112, 114, 115,116, 118, 119, 122, 124, 201, 203, 205, 207, 208, 209, 215,220, 223, and 230,\n",
    "## Test dataset record number: 100, 103, 105, 110, 113, 117, 121, 123, 200, 202, 210, 212, 213,214, 219, 221, 222, 228, 231, 232, and 234\n",
    "\n",
    "train_record_list = [101, 106, 108, 109, 112, 114, 115, 116, 118, 119, 122, 124, 201, 203, 205, 207, 208, 209, 215, 220, 223, 230]\n",
    "test_record_list = [100, 103, 105, 110, 111, 113, 117, 121, 123, 200, 202, 210, 212, 213,214, 219, 221, 222, 228, 231, 232, 233, 234]\n",
    "train_indices = [i for i in range(len(record_list)) if int(record_list[i]) in train_record_list]\n",
    "test_indices = [i for i in range(len(record_list)) if int(record_list[i]) in test_record_list]\n",
    "\n",
    "train_signals = [ecg_signals_zeropad[i][:][:] for i in train_indices]\n",
    "train_labels = [ecg_labels[i][:] for i in train_indices]\n",
    "test_signals = [ecg_signals_zeropad[i][:][:] for i in test_indices]\n",
    "test_labels = [ecg_labels[i][:] for i in test_indices]\n",
    "\n",
    "train_data = [beat_signal for patient in train_signals for beat_signal in patient]\n",
    "train_labels = [beat_label for patient in train_labels for beat_label in patient]\n",
    "test_data = [beat_signal for patient in test_signals for beat_signal in patient]\n",
    "test_labels = [beat_label for patient in test_labels for beat_label in patient]\n",
    "\n",
    "train_data = np.array(train_data)\n",
    "train_labels = np.array(train_labels)\n",
    "test_data = np.array(test_data)\n",
    "test_labels = np.array(test_labels)\n",
    "\n",
    "# TODO:\n",
    "# Save ECG 1D data sets\n",
    "with h5py.File(process_dir / 'ECG_MIT-BIH_processed_data_interpatient_360Hz.h5', 'w') as file:\n",
    "    train_set_d = file.create_dataset(\"train_data\", data=train_data)\n",
    "    train_set_l = file.create_dataset(\"train_labels\", data=train_labels)\n",
    "    test_set_d = file.create_dataset(\"test_data\", data=test_data)\n",
    "    test_set_l = file.create_dataset(\"test_labels\", data=test_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6955568-b20a-4d10-84e6-5aa1b61147e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "# Double check saved data\n",
    "# read data from h5 file\n",
    "file = h5py.File(process_dir / 'ECG_MIT-BIH_processed_data_interpatient_360Hz.h5', 'r')\n",
    "train_set_d = file['train_data'][:]\n",
    "train_set_l = file['train_labels'][:]\n",
    "test_set_d = file['test_data'][:]\n",
    "test_set_l = file['test_labels'][:]\n",
    "file.close()\n",
    "\n",
    "## Check length\n",
    "print('Train data length:', train_set_d.shape)\n",
    "print('Train label length:', train_set_l.shape)\n",
    "print(Counter(train_set_l))\n",
    "print('Test data length:', test_set_d.shape)\n",
    "print('Test label length:', test_set_l.shape)\n",
    "print(Counter(test_set_l))\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8b51224-9fba-460d-8648-c0bede360366",
   "metadata": {},
   "source": [
    "## 1.2 Intra-patient split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1916992f-d239-4bc3-8232-d07d7ccdfb32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# zero pad all heartbeat to 7 seconds\n",
    "import copy\n",
    "ecg_signals_zeropad = copy.deepcopy(ecg_signals)\n",
    "\n",
    "for i in range(len(ecg_signals)):\n",
    "    for j in range(len(ecg_signals[i])):\n",
    "\n",
    "        # pad to 256\n",
    "        zero_pad_max_length = 1000\n",
    "        pad_length = zero_pad_max_length - len(ecg_signals[i][j])\n",
    "        if pad_length <=0:\n",
    "            zero_pad_data = ecg_signals[i][j][:zero_pad_max_length]\n",
    "        else:\n",
    "            zero_pad_data = np.pad(ecg_signals[i][j], (0,pad_length),'constant')\n",
    "        ecg_signals_zeropad[i][j] = zero_pad_data\n",
    "        \n",
    "\n",
    "## plot a zero padded signal\n",
    "ecg_sample = ecg_signals_zeropad[10][100]\n",
    "\n",
    "# sampling frequency 125 Hz\n",
    "time_in_second = np.arange(len(ecg_sample)) / 360.\n",
    "\n",
    "plt.plot(time_in_second, ecg_sample, c='b', label='360 Hz')\n",
    "plt.xlabel('time in second')\n",
    "plt.ylabel('ECG value in mV')\n",
    "plt.title('Zero padded ECG signal for one beat')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "plt.clf()\n",
    "\n",
    "\n",
    "# ecg_signals_zeropad: 48 patients x number_of_beats x 2520\n",
    "# ecg_labels: 48 patients x number_of_beats\n",
    "\n",
    "## Split train and test datasets 80/20 stratified by class type\n",
    "from sklearn.model_selection import StratifiedShuffleSplit\n",
    "\n",
    "ecg_signals_zeropad_flatten = np.array([x for xs in ecg_signals_zeropad for x in xs])\n",
    "ecg_labels_flatten = np.array([x for xs in ecg_labels for x in xs])\n",
    "\n",
    "print(len(ecg_signals_zeropad_flatten), len(ecg_labels_flatten))\n",
    "\n",
    "sss = StratifiedShuffleSplit(n_splits=5, test_size=0.2, random_state=0)\n",
    "for train_index, test_index in sss.split(ecg_signals_zeropad_flatten, ecg_labels_flatten):\n",
    "    train_signals, test_signals = ecg_signals_zeropad_flatten[train_index], ecg_signals_zeropad_flatten[test_index]\n",
    "    train_labels, test_labels = ecg_labels_flatten[train_index], ecg_labels_flatten[test_index]\n",
    "    \n",
    "train_data = np.array(train_signals)\n",
    "train_labels = np.array(train_labels)\n",
    "test_data = np.array(test_signals)\n",
    "test_labels = np.array(test_labels)\n",
    "\n",
    "# Save ECG 1D data sets\n",
    "with h5py.File(process_dir / 'ECG_MIT-BIH_processed_data_intrapatient_360Hz.h5', 'w') as file:\n",
    "    train_set_d = file.create_dataset(\"train_data\", data=train_data)\n",
    "    train_set_l = file.create_dataset(\"train_labels\", data=train_labels)\n",
    "    test_set_d = file.create_dataset(\"test_data\", data=test_data)\n",
    "    test_set_l = file.create_dataset(\"test_labels\", data=test_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b5e2aa7-ce79-42d9-96ed-0068886244fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "# Double check saved data\n",
    "# read data from h5 file\n",
    "file = h5py.File(process_dir / 'ECG_MIT-BIH_processed_data_intrapatient_360Hz.h5', 'r')\n",
    "train_set_d = file['train_data'][:]\n",
    "train_set_l = file['train_labels'][:]\n",
    "test_set_d = file['test_data'][:]\n",
    "test_set_l = file['test_labels'][:]\n",
    "file.close()\n",
    "\n",
    "## Check length\n",
    "print('Train data length:', train_set_d.shape)\n",
    "print('Train label length:', train_set_l.shape)\n",
    "print(Counter(train_set_l))\n",
    "print('Test data length:', test_set_d.shape)\n",
    "print('Test label length:', test_set_l.shape)\n",
    "print(Counter(test_set_l))\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53646a6c-5724-4e5b-b77b-0a92e2ff856a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
